# -*- coding:utf-8 -*-  
# Shellcode module: manage embedded shellcodes in ROPGenerator 
import ropgenerator.Architecture as Arch
from ropgenerator.IO import error, string_special, string_bold, string_ropg, banner, notify 
import os
from prompt_toolkit import prompt
import sys
import string
import subprocess

OPTION_LIST = "--list"
OPTION_LIST_SHORT = "-l"

OPTION_ADD = "--add"
OPTION_ADD_SHORT = "-a"

OPTION_REMOVE = "--remove"
OPTION_REMOVE_SHORT = "-r"

OPTION_HELP = "--help"
OPTION_HELP_SHORT = "-h"

CMD_SHELLCODE_HELP =  banner([string_bold("'shellcode' command"),\
                    string_special("(Manage shellcodes for your exploits)")])
CMD_SHELLCODE_HELP += "\n\n\t"+string_bold("Usage")+\
"\n\t\tshellcode [OPTION] <arch>"
CMD_SHELLCODE_HELP += "\n\n\t"+string_bold("Options")+":"
CMD_SHELLCODE_HELP += "\n\t\t"+string_special(OPTION_LIST_SHORT)+","+string_special(OPTION_LIST)+"\tList available shellcodes"
CMD_SHELLCODE_HELP += "\n\t\t"+string_special(OPTION_ADD_SHORT)+","+string_special(OPTION_ADD)+"\tAdd a shellcode"
CMD_SHELLCODE_HELP += "\n\t\t"+string_special(OPTION_REMOVE_SHORT)+","+string_special(OPTION_REMOVE)+"\tRemove a previously added shellcode"
CMD_SHELLCODE_HELP += "\n\t\t"+string_special(OPTION_HELP_SHORT)+","+string_special(OPTION_HELP)+"\tShow this help"
CMD_SHELLCODE_HELP += "\n\n\t"+string_bold("Supported architectures")+": "+','.join([string_special(arch) for arch in Arch.available])

def print_help():
    print(CMD_SHELLCODE_HELP)
    
def shellcode(args):
    # Parsing arguments
    if( not args):
        print_help()
        return 
    if( args[0] in [OPTION_LIST, OPTION_LIST_SHORT]):
        func = list_shellcodes
    elif( args[0] in [OPTION_ADD, OPTION_ADD_SHORT] ):
        func = add
    elif( args[0] in [OPTION_REMOVE, OPTION_REMOVE_SHORT] ):
        func = remove
    elif( args[0] in [OPTION_HELP, OPTION_HELP_SHORT]):
        print_help()
        return 
    else:
        error("Error. Unknown option '{}'".format(args[0]))
        return 
        
    if( len(args) == 1 ):
        error("Architecture missing")
        print(string_bold("\tSupported architectures")+": "+\
            ','.join([string_special(arch) for arch in Arch.available]))
    elif( not args[1] in Arch.available ):
        error("Architecture {} not supported".format(string_special(args[1])))
        print(string_bold("\tSupported architectures")+": "+\
            ','.join([string_special(arch) for arch in Arch.available]))
    else:
        func(args[1])
    return 


def add(arch):

    print(banner([string_bold('Adding a new shellcode')]))
    
    shellcode = ''
    ok = False
    while( not ok ):
        sys.stdout.write('\t'+string_ropg('> ')+'Enter your shellcode as a string in hex format:\n')
        shellcode_input = prompt(u"        > ")
        try:
            shellcode = shellcode_input.replace('\\x','').decode('hex')
            ok = True
        except:
            ok = False
        if( not ok ):
            print(string_special("\tError. Your input is in wrong format or invalid"))

    sys.stdout.write('\t'+string_ropg('> ')+'Enter short shellcode name/description:\n')
    info = ""
    while( not info ):
        info = prompt(u"        > ")
    info =  filter( lambda x: x in set(string.printable), info)
    add_shellcode(arch, shellcode, info)


def remove(arch):
    if(not shellcodes[arch]):
        error("No shellcodes to remove for architecture " + arch)
        return 
    
    list_shellcodes(arch)
    print("")
    
    choice = ''
    ok = False
    sys.stdout.write('\t'+string_ropg('> ')+'Select a shellcode to remove:\n')
    while( not ok ):
        choice_input = prompt(u"        > ")
        try:
            choice = int(choice_input)
            ok = remove_shellcode(arch, choice)
            if( not ok ):
                print(string_special("\tError. Invalid choice"))
        except:
            ok = False

    print("")
    notify('Shellcode removed')

def select_shellcode(arch):
    """
    Returns shellcode pair (shellcode, description)
    """
    if(not shellcodes[arch]):
        error("No shellcodes available for architecture " + arch)
        return (None,None)
    
    list_shellcodes(arch)
    print("")
    
    choice = None
    ok = False
    sys.stdout.write('\t'+string_ropg('> ')+'Select a shellcode:\n')
    while( not ok ):
        choice_input = prompt(u"        > ")
        try:
            choice = int(choice_input)
            if( choice >= 1 and choice <= len(shellcodes[arch])):
                ok = True
            else:
                print(string_special("\tError. Invalid choice"))
        except:
            ok = False
    
    return  shellcodes[arch][choice-1]

##################################
# LOW LEVEL SHELLCODE MANAGMENT ##
##################################

#############
# FUNCTIONS #
#############

# Read shellcodes from a file 
def read_shellcodes(filename):
    res = [] 
    try:
        f = open(filename, 'r')
        while(True):
            l1 = f.readline()[:-1]
            l2 = f.readline()[:-1]
            if( not l2 ):
                break
            res.append([l1.decode('hex'), l2])
        f.close()
    except:
        with open(os.devnull, 'wb') as devnull:
            subprocess.call(['mkdir', DIRECTORY], stdout=devnull, stderr=devnull)
        res = []
    return res


# Save shellcodes in a file
def write_shellcodes(filename, data):
    try:
        f = open(filename, 'w')
    except:
        error("Couldn't save shellcodes: file {} can't be opened".format(filename))
        return
    for (shellcode, info) in data:
        f.write(shellcode.encode('hex')+'\n'+info+'\n')
    f.close()
    

# Save all shellcodes
def save_shellcodes():
    write_shellcodes(SHELLCODES_X86, shellcodes[Arch.ArchX86.name])
    write_shellcodes(SHELLCODES_X64, shellcodes[Arch.ArchX64.name])

def short_shellcode(raw):
        res =  "'\\x" + "\\x".join(["%02x"%ord(c) for c in raw]) + "'"
        if( len(res) > 50 ):
            res = res[:46] + "...'"
        return res

def pack_shellcode(raw):
        tmp = "\\x" + "\\x".join(["%02x"%ord(c) for c in raw])
        res = '\t'+tmp[:52]
        tmp = tmp[52:]
        while( tmp ):
            res += '\n\t'+tmp[:52]
            tmp = tmp[52:]
        return res

def show_shellcode():
    print(string_bold('\n\t-------------------------------'))
    print(string_bold("\tSelected shellcode - arch " + string_special(arch)))
    print(string_bold('\t-------------------------------'))
    (shellcode, info) = selected(arch)
    print("\n\t{}\n\n{} - {} bytes".format( info, \
            string_special(pack_shellcode(shellcode)), str(len(shellcode))))



# Print all shellcodes for an arch
def list_shellcodes(arch):
    global shellcodes

    if( arch not in Arch.available ):
        error("Error. Architecture {} is not supported".format(arch))
        return 
    if(not shellcodes[arch]):
        error("No shellcodes available for architecture " + arch)
        return 

    print(banner([string_bold("Available shellcodes for arch " + string_special(arch))]))
    i = 0
    for shellcode in shellcodes[arch]:
        i = i + 1
        number = "({})".format(string_bold(str(i)))
        print("\n\t{} {}\n\t{} - {} bytes".format(number, shellcode[1], \
            string_special(short_shellcode(shellcode[0])), str(len(shellcode[0]))))

# Add a shellcode to an arch
def add_shellcode(arch, shellcode, description):
    global shellcodes
    shellcodes[arch].append((shellcode, description))
    save_shellcodes()

# Remove a shellcode for arch
def remove_shellcode(arch, number):
    global shellcodes
    if( number < 1 or number > len(shellcodes[arch])):
        print(string_special("\tInvalid shellcode number\n"))
        return False
    del shellcodes[arch][number-1]
    save_shellcodes()
    return True


# Shellcodes files
DIRECTORY = os.path.expanduser('~')+"/.ROPGenerator/"
SHELLCODES_X86 = DIRECTORY + "shellcodes_X86"
SHELLCODES_X64 = DIRECTORY + "shellcodes_X64"

# A shellcode is stored as a pair (shellcode, description)
# Keys are the architectures names (arch.name)
shellcodes = dict()
shellcodes[Arch.ArchX86.name] = read_shellcodes(SHELLCODES_X86)
shellcodes[Arch.ArchX64.name] = read_shellcodes(SHELLCODES_X64)

